// Copyright 2021 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#ifndef IREE_COMPILER_BINDINGS_TFLITE_TRANSFORMS_PASSES_H_
#define IREE_COMPILER_BINDINGS_TFLITE_TRANSFORMS_PASSES_H_

#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Support/LLVM.h"

namespace mlir {
namespace iree_compiler {
namespace IREE {
namespace TFLite {

//===----------------------------------------------------------------------===//
// Helpers
//===----------------------------------------------------------------------===//

// Adds a set of passes to the given pass manager that setup a module for use
// with the IREE TFLite runtime bindings.
void buildTransformPassPipeline(OpPassManager &passManager);

void registerTransformPassPipeline();

//===----------------------------------------------------------------------===//
// TFLite bindings support
//===----------------------------------------------------------------------===//

// Wraps all model entry points in a function that is compatible with the
// expected invocation semantics of the IREE TFLite bindings.
std::unique_ptr<OperationPass<ModuleOp>> createWrapEntryPointsPass();

// Materialize the functions required by the runtime bindings to manipulate
// the program state (like _resize_input_shape, etc).
std::unique_ptr<OperationPass<ModuleOp>> createMaterializeShapeSupportPass();

//===----------------------------------------------------------------------===//
// Register all Passes
//===----------------------------------------------------------------------===//

inline void registerPasses() {
  createWrapEntryPointsPass();
  createMaterializeShapeSupportPass();
}

}  // namespace TFLite
}  // namespace IREE
}  // namespace iree_compiler
}  // namespace mlir

#endif  // IREE_COMPILER_BINDINGS_TFLITE_TRANSFORMS_PASSES_H_
